# Chinese Military News
import os
import random

import torch
from torch.utils.data import Dataset
from utils.utils import read_from_txt


class CMNDataset(Dataset):
    def __init__(self, config, split, rebuild_vocab=False):
        self.config = config
        if split == 'train':
            self.texts, self.text_labels = read_from_txt(self.config.train_path)
        if split == 'val':
            self.texts, self.text_labels = read_from_txt(self.config.val_path)
        if split == 'test':
            self.texts, self.text_labels = read_from_txt(self.config.test_path)
        
        self.text_labels = torch.LongTensor(self.text_labels).to(self.config.device)

    def __len__(self):
        return len(self.text_labels)

    def __getitem__(self, idx):
        text = self.texts[idx]
        label = self.text_labels[idx]
        return text, label
